#!/bin/bash
#SBATCH -p kshdtest       # 分区名称
#SBATCH --gres=dcu:1      # 请求1个DCU资源
#SBATCH -n 1            # MPI进程数
#SBATCH -c 8            # CPU核心数 (减少核心数以快速分配)
#SBATCH -t 01:00:00     # 最大运行时间 (增加到1小时以适应迭代)
#SBATCH -J AdaptiveMatMulTest # 作业名称 (更新)
#SBATCH -o %x_%j.out    # 标准输出文件
#SBATCH -e %x_%j.err    # 错误输出文件

# --- Adaptive Test Script ---

echo "=== Adaptive Test Starting: $(date) ===="
echo "Job ID: $SLURM_JOB_ID"
echo "Node: $SLURM_JOB_NODELIST"
echo "Working Directory: $(pwd)"

module purge
module load compiler/rocm/dtk-23.10
module load compiler/gcc/9.3.0
module load mathlib/blas/3.5.0/intel

echo "=== Environment ===="
echo "HIP_PATH: ${HIP_PATH}"
echo "ROCM_PATH: ${ROCM_PATH}"
echo ""
# 2. Compile Code
echo "--- Compiling Code ---"
make clean
# Make sure the executable name matches the one used below
EXECUTABLE="./matmul" # Changed from matrix_multiply to match Makefile target
make -j HIPCC=hipcc # This builds the default target which is 'matmul'

if [ $? -ne 0 ]; then
    echo "!!! Compilation FAILED !!!"
    exit 1
fi
echo "Compilation successful."
echo ""

# 3. Run Adaptive Tests
echo "--- Running Adaptive Performance Tests ---"
LOG_FILE="adaptive_test_output.log"
SUMMARY_FILE="adaptive_summary.log"
echo "Test output will be logged to $LOG_FILE"
echo "Summary will be logged to $SUMMARY_FILE"
echo "" > $LOG_FILE # Clear log file
echo "Size,Best_Hybrid_Time_ms,Best_CPU_Ratio,Best_Iteration,Avg_CPU_Time_ms,Avg_DCU_Time_ms" > $SUMMARY_FILE # CSV Header

SIZES=(512 1024 2048 4096 8192)
# ADAPTIVE_ITERATIONS=15 # Total iterations for adaptive hybrid mode per size
ADAPTIVE_ITERATIONS=10 # Reduced for quicker test
# RUNS_PER_MODE=3      # Removed - run each mode only once per iteration
INITIAL_HYBRID_RATIO=0.5 # Starting ratio for hybrid mode

for SIZE in ${SIZES[@]}; do
    M=$SIZE
    K=$SIZE
    N=$SIZE

    echo "--- Testing Size: ${M}x${K}x${N} ---" | tee -a $LOG_FILE

    # --- Reset state for this size ---
    declare -a cpu_times_ms
    declare -a dcu_times_ms
    declare -a hybrid_times_ms # Stores avg hybrid time per iteration
    declare -a hybrid_ratios_used # Stores ratio used for hybrid per iteration

    cumulative_cpu_time_ms=0.0
    cumulative_dcu_time_ms=0.0
    total_cpu_runs=0
    total_dcu_runs=0

    best_hybrid_avg_time_ms=-1.0
    best_hybrid_ratio=-1.0
    best_hybrid_iteration=-1

    current_cpu_ratio=$INITIAL_HYBRID_RATIO

    # --- Adaptive Iteration Loop for Hybrid Mode ---
    for iter in $(seq 1 $ADAPTIVE_ITERATIONS); do
        echo -e "
--- Size $SIZE - Iteration $iter / $ADAPTIVE_ITERATIONS ---" | tee -a $LOG_FILE

        # --- Run Pure CPU Mode (single run) ---
        echo "Running CPU Mode (1 time)..." | tee -a $LOG_FILE
        # iter_cpu_time_sum=0.0 # Removed
        # for run in $(seq 1 $RUNS_PER_MODE); do # Removed inner loop
        # Use simple echo piping for dimensions, matching main.cpp expectation
        output_cpu=$(echo "$M $K $N" | $EXECUTABLE -c)
        time_ms_cpu=$(echo "$output_cpu" | grep "RESULT_TIME_MS:" | awk '{print $2}')
        if [[ -z "$time_ms_cpu" || ! "$time_ms_cpu" =~ ^[0-9]+([.][0-9]+)?$ ]]; then
            echo "错误：无法从 CPU 运行 $iter 的输出中解析有效时间！ Output:" | tee -a $LOG_FILE
            echo "$output_cpu" | tee -a $LOG_FILE
            exit 1
        fi
        # iter_cpu_time_sum=$(echo "$iter_cpu_time_sum + $time_ms" | bc) # Removed
        cpu_times_ms+=($time_ms_cpu) # Still useful to keep array of individual times if needed later
        # done # Removed inner loop end
        # if [[ $RUNS_PER_MODE -gt 0 ]]; then # Removed condition
        cumulative_cpu_time_ms=$(echo "$cumulative_cpu_time_ms + $time_ms_cpu" | bc)
        total_cpu_runs=$(($total_cpu_runs + 1)) # Increment by 1
        echo "CPU Time (this iter): $time_ms_cpu ms" | tee -a $LOG_FILE
        # fi # Removed condition end


        # --- Run Pure DCU Mode (single run) ---
        echo "Running DCU Mode (1 time)..." | tee -a $LOG_FILE
        # iter_dcu_time_sum=0.0 # Removed
        # for run in $(seq 1 $RUNS_PER_MODE); do # Removed inner loop
        output_dcu=$(echo "$M $K $N" | $EXECUTABLE -d)
        time_ms_dcu=$(echo "$output_dcu" | grep "RESULT_TIME_MS:" | awk '{print $2}')
        if [[ -z "$time_ms_dcu" || ! "$time_ms_dcu" =~ ^[0-9]+([.][0-9]+)?$ ]]; then
            echo "错误：无法从 DCU 运行 $iter 的输出中解析有效时间！ Output:" | tee -a $LOG_FILE
            echo "$output_dcu" | tee -a $LOG_FILE
            exit 1
        fi
        # iter_dcu_time_sum=$(echo "$iter_dcu_time_sum + $time_ms" | bc) # Removed
        dcu_times_ms+=($time_ms_dcu)
        # done # Removed inner loop end
        # if [[ $RUNS_PER_MODE -gt 0 ]]; then # Removed condition
        cumulative_dcu_time_ms=$(echo "$cumulative_dcu_time_ms + $time_ms_dcu" | bc)
        total_dcu_runs=$(($total_dcu_runs + 1)) # Increment by 1
        echo "DCU Time (this iter): $time_ms_dcu ms" | tee -a $LOG_FILE
        # fi # Removed condition end

        # --- Calculate cumulative averages and next cpuRatio (from iter 2) --- 
        # This part remains largely the same as it uses cumulative values
        if [[ $iter -gt 1 ]]; then
            if [[ $total_cpu_runs -gt 0 && $total_dcu_runs -gt 0 ]]; then
                t_cpu_avg=$(echo "scale=6; $cumulative_cpu_time_ms / $total_cpu_runs" | bc)
                t_dcu_avg=$(echo "scale=6; $cumulative_dcu_time_ms / $total_dcu_runs" | bc)
                echo "Cumulative Avg: CPU=$t_cpu_avg ms, DCU=$t_dcu_avg ms" | tee -a $LOG_FILE

                denominator=$(echo "$t_cpu_avg + $t_dcu_avg" | bc)
                if (( $(echo "$denominator > 0.000001" | bc -l) )); then
                    current_cpu_ratio=$(echo "scale=5; $t_dcu_avg / $denominator" | bc)
                    if (( $(echo "$current_cpu_ratio < 0" | bc -l) )); then current_cpu_ratio=0.0; fi
                    if (( $(echo "$current_cpu_ratio > 1" | bc -l) )); then current_cpu_ratio=1.0; fi
                     echo "Calculated next Hybrid CPU Ratio: $current_cpu_ratio" | tee -a $LOG_FILE
                else
                    echo "Warning: Sum of avg CPU and DCU times is zero or too small. Using default ratio 0.5" | tee -a $LOG_FILE
                    current_cpu_ratio=0.5
                fi
            else
                 echo "Warning: Not enough historical data for CPU/DCU runs. Using previous ratio: $current_cpu_ratio" | tee -a $LOG_FILE
            fi
        else
             echo "Iteration 1, using initial Hybrid CPU Ratio: $current_cpu_ratio" | tee -a $LOG_FILE
        fi
        hybrid_ratios_used+=($current_cpu_ratio)

        # --- Run Hybrid Mode (single run) ---
        echo "Running Hybrid Mode (CPU Ratio: $current_cpu_ratio, 1 time)..." | tee -a $LOG_FILE
        # iter_hybrid_time_sum=0.0 # Removed
        # Make sure executable name matches the compiled one
        # for run in $(seq 1 $RUNS_PER_MODE); do # Removed inner loop
        output_hybrid=$(echo "$M $K $N" | $EXECUTABLE -h --force-cpu-ratio $current_cpu_ratio)
        time_ms_hybrid=$(echo "$output_hybrid" | grep "RESULT_TIME_MS:" | awk '{print $2}')
         if [[ -z "$time_ms_hybrid" || ! "$time_ms_hybrid" =~ ^[0-9]+([.][0-9]+)?$ ]]; then
            echo "错误：无法从混合模式运行 $iter 的输出中解析有效时间！ Output:" | tee -a $LOG_FILE
             echo "$output_hybrid" | tee -a $LOG_FILE
            exit 1
        fi
        # iter_hybrid_time_sum=$(echo "$iter_hybrid_time_sum + $time_ms" | bc) # Removed
        # done # Removed inner loop end
         hybrid_times_ms+=($time_ms_hybrid) # Store the single time for this iteration
        echo "Hybrid Time (this iter): $time_ms_hybrid ms" | tee -a $LOG_FILE

        # --- Update best result for this SIZE --- 
        # Compare the current single hybrid time with the best single time found so far
        is_first_best=0
        [[ $(echo "$best_hybrid_avg_time_ms == -1.0" | bc -l) -eq 1 ]] && is_first_best=1

        is_better=0
        [[ $is_first_best -eq 0 && $(echo "$time_ms_hybrid < $best_hybrid_avg_time_ms" | bc -l) -eq 1 ]] && is_better=1

        if [[ $is_first_best -eq 1 || $is_better -eq 1 ]]; then
             best_hybrid_avg_time_ms=$time_ms_hybrid # Store the new best single time
             best_hybrid_ratio=$current_cpu_ratio
             best_hybrid_iteration=$iter
             echo "  *** New best hybrid result for size $SIZE found! (Iter $iter) ***" | tee -a $LOG_FILE
        fi
        # fi # Removed condition end
    done # End of adaptive iterations for this size

    # --- Log summary for this size --- 
    # The summary log remains the same, reporting the best single hybrid time found
    final_avg_cpu="N/A"
    final_avg_dcu="N/A"
     if [[ $total_cpu_runs -gt 0 ]]; then final_avg_cpu=$(echo "scale=3; $cumulative_cpu_time_ms / $total_cpu_runs" | bc); fi
     if [[ $total_dcu_runs -gt 0 ]]; then final_avg_dcu=$(echo "scale=3; $cumulative_dcu_time_ms / $total_dcu_runs" | bc); fi

    echo "${SIZE}x${SIZE},${best_hybrid_avg_time_ms},${best_hybrid_ratio},${best_hybrid_iteration},${final_avg_cpu},${final_avg_dcu}" >> $SUMMARY_FILE
    echo "--- Finished Size $SIZE ---" | tee -a $LOG_FILE
    echo "" | tee -a $LOG_FILE

done # End of loop over SIZES

# 4. Test Python Environment (Keep as is)
echo "--- Testing Python Environment ---" | tee -a $LOG_FILE
python3 plot_performance.py >> $LOG_FILE 2>&1
if [ $? -eq 0 ]; then
    echo "Python script execution attempt finished (check log for success/errors)." | tee -a $LOG_FILE
else
    echo "!!! Python script execution FAILED (python3 or modules not found) !!!" | tee -a $LOG_FILE
fi
echo ""

# 5. Finish
echo "=== Adaptive Test Finished: $(date) ====" | tee -a $LOG_FILE
echo "Check $LOG_FILE for full test output."
echo "Check $SUMMARY_FILE for best hybrid results per size."
echo "Check Slurm output files (%x_%j.out, %x_%j.err) for other messages."

exit 0 